package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 对于生产环境，请根据实际情况调整此函数
	},
}

// 定义一个全局变量来存储所有的连接
var clients = make(map[*websocket.Conn]bool) // 存储所有连接
var broadcast = make(chan []byte)            // 广播通道
var mutex sync.Mutex                         // 互斥锁，用于保护clients

func wsHandler(c *gin.Context) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("Failed to set websocket upgrade: %v", err)
		return
	}
	defer func() {
		mutex.Lock()
		delete(clients, conn)
		mutex.Unlock()
		conn.Close()
	}()

	mutex.Lock()
	clients[conn] = true
	mutex.Unlock()

	for {
		messageType, message, err := conn.ReadMessage()
		fmt.Println(string(message))
		if err != nil {
			log.Printf("Error reading message: %v", err)
			break
		}

		// 将接收到的消息发送到广播通道
		if messageType == websocket.TextMessage {
			broadcast <- message
		}
	}
}

func broadcaster() {
	for {
		msg := <-broadcast

		// 使用互斥锁遍历所有连接并发送消息
		mutex.Lock()
		for client := range clients {
			err := client.WriteMessage(websocket.TextMessage, msg)
			if err != nil {
				log.Printf("Error writing to client: %v", err)
				client.Close()
				delete(clients, client)
			}
		}
		mutex.Unlock()
	}
}

func main() {
	go broadcaster() // 启动广播器

	r := gin.Default()
	r.GET("/ws", wsHandler)
	r.Run(":8080") // 监听并在 0.0.0.0:8080 上启动服务
}
